package io.openbas.migration;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.Statement;
import org.flywaydb.core.api.migration.BaseJavaMigration;
import org.flywaydb.core.api.migration.Context;
import org.springframework.stereotype.Component;

@Component
public class V4_06__Fix_vulnerability_expectations extends BaseJavaMigration {

  @Override
  public void migrate(Context context) throws Exception {
    Statement statement = context.getConnection().createStatement();
    ObjectMapper mapper = new ObjectMapper();

    // Delete incorrect collector
    statement.executeUpdate(
        "DELETE FROM collectors WHERE collector_id = 'acab8214-0379-448a-a575-05e9d934eadd'");

    // Fix injectors_contracts
    ResultSet injectorContracts =
        statement.executeQuery(
            "SELECT injector_contract_content, injector_contract_id FROM injectors_contracts WHERE injector_contract_payload IS NOT NULL");

    PreparedStatement updateContracts =
        context
            .getConnection()
            .prepareStatement(
                "UPDATE injectors_contracts SET injector_contract_content = ? WHERE injector_contract_id = ?");

    while (injectorContracts.next()) {
      String content = injectorContracts.getString("injector_contract_content");
      String contractId = injectorContracts.getString("injector_contract_id");

      ObjectNode contractContent = (ObjectNode) mapper.readTree(content);
      boolean modified = false;

      if (contractContent != null && contractContent.has("fields")) {
        ArrayNode fields = (ArrayNode) contractContent.get("fields");

        for (int i = 0; i < fields.size(); i++) {
          ObjectNode field = (ObjectNode) fields.get(i);

          if ("expectations".equals(field.path("key").asText())) {
            ArrayNode expectations = (ArrayNode) field.path("predefinedExpectations");

            for (JsonNode expectation : expectations) {
              if ("VULNERABILITY".equals(expectation.path("expectation_type").asText())) {
                String name = expectation.path("expectation_name").asText();
                if (!"Not vulnerable".equals(name)) {
                  ((ObjectNode) expectation).put("expectation_name", "Not vulnerable");
                  modified = true;
                }
              }
            }
          }
        }

        if (modified) {
          contractContent.set("fields", fields);
          String updatedContent = mapper.writeValueAsString(contractContent);
          updateContracts.setString(1, updatedContent);
          updateContracts.setString(2, contractId);
          updateContracts.addBatch();
        }
      }
    }
    updateContracts.executeBatch();

    // Fix injects.inject_content
    ResultSet injects =
        statement.executeQuery(
            "SELECT inject_id, inject_content FROM injects WHERE inject_content IS NOT NULL");

    PreparedStatement updateInjects =
        context
            .getConnection()
            .prepareStatement("UPDATE injects SET inject_content = ? WHERE inject_id = ?");

    while (injects.next()) {
      String content = injects.getString("inject_content");
      String injectId = injects.getString("inject_id");

      JsonNode root = mapper.readTree(content);
      if (root != null && root.isObject()) {

        ObjectNode injectContent = (ObjectNode) root;
        boolean modified = false;

        if (injectContent.has("expectations") && injectContent.get("expectations").isArray()) {
          ArrayNode expectations = (ArrayNode) injectContent.get("expectations");

          for (JsonNode expectation : expectations) {
            if ("VULNERABILITY".equals(expectation.path("expectation_type").asText())) {
              String name = expectation.path("expectation_name").asText();
              if (!"Not vulnerable".equals(name)) {
                ((ObjectNode) expectation).put("expectation_name", "Not vulnerable");
                modified = true;
              }
            }
          }

          if (modified) {
            String updatedContent = mapper.writeValueAsString(injectContent);
            updateInjects.setString(1, updatedContent);
            updateInjects.setString(2, injectId);
            updateInjects.addBatch();
          }
        }
      }
    }

    updateInjects.executeBatch();
  }
}
